{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install dependencies for Colab\\n\n",
    "%pip -q install -U pip\n",
    "%pip -q install llama-index-retrievers-superlinked"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example: Superlinked + LlamaIndex custom retriever (Steam games)\n",
    "# This notebook mirrors examples/steam_games_example.py\n",
    "\n",
    "import argparse\n",
    "from typing import List, Optional\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "import superlinked.framework as sl\n",
    "from llama_index.retrievers.superlinked import SuperlinkedRetriever\n",
    "\n",
    "try:\n",
    "    from llama_index.core.query_engine import RetrieverQueryEngine\n",
    "    from llama_index.core.response_synthesizers import get_response_synthesizer\n",
    "except Exception:\n",
    "    RetrieverQueryEngine = None  # type: ignore\n",
    "    get_response_synthesizer = None  # type: ignore\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_dataframe(csv_path: Optional[str]) -> pd.DataFrame:\n",
    "    if csv_path:\n",
    "        df = pd.read_csv(csv_path)\n",
    "    else:\n",
    "        df = pd.DataFrame(\n",
    "            [\n",
    "                {\n",
    "                    \"game_number\": 1,\n",
    "                    \"name\": \"Star Tactics\",\n",
    "                    \"desc_snippet\": \"Turn-based strategy in deep space.\",\n",
    "                    \"game_details\": \"Tactical combat, fleet management\",\n",
    "                    \"languages\": \"en\",\n",
    "                    \"genre\": \"Strategy, Sci-Fi\",\n",
    "                    \"game_description\": \"Engage in strategic battles among the stars.\",\n",
    "                    \"original_price\": 29.99,\n",
    "                    \"discount_price\": 19.99,\n",
    "                },\n",
    "                {\n",
    "                    \"game_number\": 2,\n",
    "                    \"name\": \"Wizard Party\",\n",
    "                    \"desc_snippet\": \"Co-op party game with spells.\",\n",
    "                    \"game_details\": \"Local co-op, party\",\n",
    "                    \"languages\": \"en\",\n",
    "                    \"genre\": \"Party, Casual, Magic\",\n",
    "                    \"game_description\": \"Cast spells with friends in chaotic party modes.\",\n",
    "                    \"original_price\": 14.99,\n",
    "                    \"discount_price\": 9.99,\n",
    "                },\n",
    "            ]\n",
    "        )\n",
    "\n",
    "    required = [\n",
    "        \"game_number\",\n",
    "        \"name\",\n",
    "        \"desc_snippet\",\n",
    "        \"game_details\",\n",
    "        \"languages\",\n",
    "        \"genre\",\n",
    "        \"game_description\",\n",
    "        \"original_price\",\n",
    "        \"discount_price\",\n",
    "    ]\n",
    "    missing = [c for c in required if c not in df.columns]\n",
    "    if missing:\n",
    "        raise ValueError(f\"Missing required columns: {missing}\")\n",
    "\n",
    "    df[\"combined_text\"] = (\n",
    "        df[\"name\"].astype(str)\n",
    "        + \" \"\n",
    "        + df[\"desc_snippet\"].astype(str)\n",
    "        + \" \"\n",
    "        + df[\"genre\"].astype(str)\n",
    "        + \" \"\n",
    "        + df[\"game_details\"].astype(str)\n",
    "        + \" \"\n",
    "        + df[\"game_description\"].astype(str)\n",
    "    )\n",
    "    return df\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_superlinked_app(df: pd.DataFrame):\n",
    "    class GameSchema(sl.Schema):\n",
    "        id: sl.IdField\n",
    "        name: sl.String\n",
    "        desc_snippet: sl.String\n",
    "        game_details: sl.String\n",
    "        languages: sl.String\n",
    "        genre: sl.String\n",
    "        game_description: sl.String\n",
    "        original_price: sl.Float\n",
    "        discount_price: sl.Float\n",
    "        combined_text: sl.String\n",
    "\n",
    "    game = GameSchema()\n",
    "\n",
    "    text_space = sl.TextSimilaritySpace(\n",
    "        text=game.combined_text,\n",
    "        model=\"sentence-transformers/all-mpnet-base-v2\",\n",
    "    )\n",
    "    index = sl.Index([text_space])\n",
    "\n",
    "    parser = sl.DataFrameParser(\n",
    "        game,\n",
    "        mapping={\n",
    "            game.id: \"game_number\",\n",
    "            game.name: \"name\",\n",
    "            game.desc_snippet: \"desc_snippet\",\n",
    "            game.game_details: \"game_details\",\n",
    "            game.languages: \"languages\",\n",
    "            game.genre: \"genre\",\n",
    "            game.game_description: \"game_description\",\n",
    "            game.original_price: \"original_price\",\n",
    "            game.discount_price: \"discount_price\",\n",
    "            game.combined_text: \"combined_text\",\n",
    "        },\n",
    "    )\n",
    "\n",
    "    source = sl.InMemorySource(schema=game, parser=parser)\n",
    "    executor = sl.InMemoryExecutor(sources=[source], indices=[index])\n",
    "    app = executor.run()\n",
    "\n",
    "    source.put([df])\n",
    "\n",
    "    query = (\n",
    "        sl.Query(index)\n",
    "        .find(game)\n",
    "        .similar(text_space, sl.Param(\"query_text\"))\n",
    "        .select(\n",
    "            [\n",
    "                game.id,\n",
    "                game.name,\n",
    "                game.desc_snippet,\n",
    "                game.game_details,\n",
    "                game.languages,\n",
    "                game.genre,\n",
    "                game.game_description,\n",
    "                game.original_price,\n",
    "                game.discount_price,\n",
    "            ]\n",
    "        )\n",
    "    )\n",
    "\n",
    "    return app, query, game\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_demo(csv_path: Optional[str], top_k: int, query_text: str) -> None:\n",
    "    df = build_dataframe(csv_path)\n",
    "    app, query_descriptor, game = build_superlinked_app(df)\n",
    "\n",
    "    retriever = SuperlinkedRetriever(\n",
    "        sl_client=app,\n",
    "        sl_query=query_descriptor,\n",
    "        page_content_field=\"desc_snippet\",\n",
    "        query_text_param=\"query_text\",\n",
    "        metadata_fields=[\n",
    "            \"id\",\n",
    "            \"name\",\n",
    "            \"genre\",\n",
    "            \"game_details\",\n",
    "            \"languages\",\n",
    "            \"game_description\",\n",
    "            \"original_price\",\n",
    "            \"discount_price\",\n",
    "        ],\n",
    "        top_k=top_k,\n",
    "    )\n",
    "\n",
    "    print(f\"\\nRetrieving for: {query_text!r}\")\n",
    "    nodes = retriever.retrieve(query_text)\n",
    "    for i, nws in enumerate(nodes, 1):\n",
    "        print(f\"#{i} score={nws.score:.4f} text={nws.node.text!r}\")\n",
    "        print(f\"   metadata: {nws.node.metadata}\")\n",
    "\n",
    "    if RetrieverQueryEngine and get_response_synthesizer:\n",
    "        print(\"\\nBuilding RetrieverQueryEngine...\")\n",
    "        try:\n",
    "            engine = RetrieverQueryEngine(\n",
    "                retriever=retriever, response_synthesizer=get_response_synthesizer()\n",
    "            )\n",
    "            response = engine.query(query_text)\n",
    "            print(\"\\nEngine response:\", response)\n",
    "        except Exception as e:\n",
    "            print(\"Engine invocation failed (likely missing LLM setup):\", e)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parameters (for Colab users)\n",
    "csv_path = None  # @param {type:\"string\"}\n",
    "top_k = 3        # @param {type:\"integer\"}\n",
    "query_text = \"strategic sci-fi game\"  # @param {type:\"string\"}\n",
    "\n",
    "run_demo(csv_path, top_k, query_text)\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
